package rocket

import (
	"context"

	"github.com/apache/rocketmq-client-go/v2"
	"github.com/apache/rocketmq-client-go/v2/primitive"
	"github.com/apache/rocketmq-client-go/v2/producer"
	"github.com/zeromicro/go-zero/core/logx"

	"gitlab.vspncd.com/mic/common/tools"
)

type (
	Producer interface {
		// 开启
		Start()
		// 关闭
		Stop()
		// 发送延迟消息
		SendDelayMsg(ctx context.Context, key string, body []byte, delayLevel int) (string, error)
		// 发送消息
		SendMsg(ctx context.Context, body []byte) (string, error)
		SendMsgWithKey(ctx context.Context, key string, body []byte) (string, error)
	}

	defaultProducer struct {
		producer rocketmq.Producer
		topic    string
	}
)

func NewProducer(cfg *ProducerConfig) Producer {
	pAddr, err := primitive.NewNamesrvAddr(cfg.Addr...)
	if err != nil {
		panic(err)
	}

	opts := make([]producer.Option, 0, 5)
	opts = append(opts, producer.WithNameServer(pAddr))
	opts = append(opts, producer.WithRetry(2))
	opts = append(opts, producer.WithGroupName(cfg.Topic))
	if cfg.Auth != nil {
		opts = append(opts, producer.WithCredentials(primitive.Credentials{
			AccessKey:     cfg.Auth.AccessKey,
			SecretKey:     cfg.Auth.SecretKey,
			SecurityToken: cfg.Auth.SecurityToken,
		}))
	}

	p, err := rocketmq.NewProducer(opts...)
	if err != nil {
		panic(err)
	}

	return &defaultProducer{
		producer: p,
		topic:    cfg.Topic,
	}
}

func (d *defaultProducer) Start() {
	err := d.producer.Start()
	if err != nil {
		panic(err)
	}
}

func (d *defaultProducer) Stop() {
	_ = d.producer.Shutdown()
}

func (d *defaultProducer) SendDelayMsg(ctx context.Context, key string, body []byte, delayLevel int) (string, error) {
	msg := primitive.NewMessage(d.topic, body)
	msg.WithDelayTimeLevel(delayLevel)
	msg.WithTag(key)
	return d.send(ctx, msg)
}

func (d *defaultProducer) SendMsg(ctx context.Context, body []byte) (string, error) {
	msg := primitive.NewMessage(d.topic, body)
	return d.send(ctx, msg)
}

func (d *defaultProducer) SendMsgWithKey(ctx context.Context, key string, body []byte) (string, error) {
	msg := primitive.NewMessage(d.topic, body)
	msg.WithTag(key)
	return d.send(ctx, msg)
}

func (d *defaultProducer) send(ctx context.Context, msg *primitive.Message) (string, error) {
	var msgId string
	err := tools.DoWithRetry(func() error {
		res, e := d.producer.SendSync(ctx, msg)
		if e != nil {
			return e
		}
		msgId = res.MsgID
		return nil
	})
	if err != nil {
		logx.Error("send msg error", err)
		return "", err
	}

	return msgId, nil
}
